# Copyright (c) mask_decoderMeta Platforms, Inc. and affiliates.
# All rights reserved.
import warnings

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
from torch import nn
from torch.nn import functional as F

from typing import List, Tuple, Type

from .common import LayerNorm2d
import math


# new!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
    # type: (Tensor, float, float, float, float) -> Tensor
    r"""Fills the input Tensor with values drawn from a truncated
    normal distribution. The values are effectively drawn from the
    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
    with values outside :math:`[a, b]` redrawn until they are within
    the bounds. The method used for generating the random values works
    best when :math:`a \leq \text{mean} \leq b`.
    Args:
        tensor: an n-dimensional `torch.Tensor`
        mean: the mean of the normal distribution
        std: the standard deviation of the normal distribution
        a: the minimum cutoff value
        b: the maximum cutoff value
    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.trunc_normal_(w)
    """
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


# new!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
def _no_grad_trunc_normal_(tensor, mean, std, a, b):
    # Cut & paste from PyTorch official master until it's in a few official releases - RW
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x):
        # Computes standard normal cumulative distribution function
        return (1. + math.erf(x / math.sqrt(2.))) / 2.

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
                      "The distribution of values may be incorrect.",
                      stacklevel=2)

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        # Get upper and lower cdf values
        l = norm_cdf((a - mean) / std)
        u = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [l, u], then translate to
        # [2l-1, 2u-1].
        tensor.uniform_(2 * l - 1, 2 * u - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


class MaskDecoder(nn.Module):
    def __init__(
            self,
            *,
            transformer_dim: int,
            transformer: nn.Module,
            num_multimask_outputs: int = 4,
            activation: Type[nn.Module] = nn.GELU,
            iou_head_depth: int = 3,
            iou_head_hidden_dim: int = 256,
    ) -> None:
        """
        Predicts masks given an image and prompt embeddings, using a
        transformer architecture.

        Arguments:
          transformer_dim (int): the channel dimension of the transformer
          transformer (nn.Module): the transformer used to predict masks
          num_multimask_outputs (int): the number of masks to predict
            when disambiguating masks
          activation (nn.Module): the type of activation to use when
            upscaling masks
          iou_head_depth (int): the depth of the MLP used to predict
            mask quality
          iou_head_hidden_dim (int): the hidden dimension of the MLP
            used to predict mask quality
        """
        super().__init__()
        self.transformer_dim = transformer_dim
        self.transformer = transformer

        self.num_multimask_outputs = num_multimask_outputs

        self.iou_token = nn.Embedding(1, transformer_dim)  # (1,256)
        self.num_mask_tokens = num_multimask_outputs + 1
        self.mask_tokens = nn.Embedding(self.num_mask_tokens, transformer_dim)  # (4,256)

        # output_upscaling序列模型的作用是对输入进行两次上采样，并在上采样过程中进行标准化和非线性变换，最终得到上采样后的输出结果。这种上采样操作通常用于将低分辨率的特征图上采样到高分辨率，或者进行特征提取以增加模型的感知能力。
        self.output_upscaling = nn.Sequential(
            nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2),
            # 转置卷积层，用于对输入进行上采样，参数是通道数，H*W->2H*2W
            LayerNorm2d(transformer_dim // 4),  # LN标准化处理
            activation(),
            nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
            activation(),
        )
        self.output_hypernetworks_mlps = nn.ModuleList(
            [
                MLP(transformer_dim, transformer_dim, transformer_dim // 8, 3)
                for i in range(self.num_mask_tokens)
            ]
        )

        self.iou_prediction_head = MLP(
            transformer_dim, iou_head_hidden_dim, self.num_mask_tokens, iou_head_depth
        )

        # new!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!finetune 同upscaling
        self.output_upscaling_adaptor = nn.Sequential(
            nn.ConvTranspose2d(transformer_dim, transformer_dim // 4, kernel_size=2, stride=2),
            LayerNorm2d(transformer_dim // 4),
            activation(),
            nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
            activation(),
        )

        # new!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! 接在两个卷积后面
        self.cls_upscaling = nn.Sequential(
            nn.ConvTranspose2d(transformer_dim // 4, transformer_dim // 8, kernel_size=2, stride=2),
            LayerNorm2d(transformer_dim // 8),
            activation(),
            nn.Conv2d(transformer_dim // 8, transformer_dim * 5 // 8, kernel_size=7, stride=2, padding=3),
            activation(),
        )
        # self.cls_tokens_embeddings = nn.ModuleList(
        #     [nn.Embedding(4,transformer_dim)  #4-num_cls
        #      for i in range(self.num_mask_tokens)
        #      ]
        # )
        # self.cls_tokens = nn.Embedding(4,transformer_dim)

        self.apply(self._init_weights)

    # new!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!初始化模型的权重和偏置
    def _init_weights(self, m):
        if isinstance(m, nn.Linear):  # 线性层
            trunc_normal_(m.weight, std=.02)
            # 对该线性层的权重进行截断正态分布初始化，标准差为 0.02
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)  # 如果该线性层有偏置项，将偏置项初始化为常数 0
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)  # 偏置项初始化为常数 0
            nn.init.constant_(m.weight, 1.0)  # 权重初始化为常数 1.0
        elif isinstance(m, nn.Conv2d):
            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels  # 计算卷积核的输出通道数，一个卷积核对应一个输出通道特征图
            fan_out //= m.groups  # 如果卷积层使用了分组卷积，则将输出通道数除以分组数，得到每个分组的输出通道数。
            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))  # 对卷积层的权重进行正态分布初始化，均值为 0，标准差为 2.0 除以 fan_out 的平方根
            if m.bias is not None:
                m.bias.data.zero_()  # 偏置项初始化为常数 0

    def forward(
            self,
            image_embeddings: torch.Tensor,
            image_pe: torch.Tensor,
            sparse_prompt_embeddings: torch.Tensor,
            dense_prompt_embeddings: torch.Tensor,
            multimask_output: bool,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Predict masks given image and prompt embeddings.

        Arguments:
          image_embeddings (torch.Tensor): the embeddings from the image encoder
          image_pe (torch.Tensor): positional encoding with the shape of image_embeddings
          sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes
          dense_prompt_embeddings (torch.Tensor): the embeddings of the mask inputs
          multimask_output (bool): Whether to return multiple masks or a single
            mask.

        Returns:
          torch.Tensor: batched predicted masks
          torch.Tensor: batched predictions of mask quality
        """
        masks, iou_pred = self.predict_masks(
            image_embeddings=image_embeddings,
            image_pe=image_pe,
            sparse_prompt_embeddings=sparse_prompt_embeddings,
            dense_prompt_embeddings=dense_prompt_embeddings,
        )

        # Select the correct mask or masks for outptu
        if multimask_output:
            mask_slice = slice(1, None)
        else:
            mask_slice = slice(0, 1)
        masks = masks[:, mask_slice, :, :, :]
        iou_pred = iou_pred[:, mask_slice]

        # Prepare output
        return masks, iou_pred

    def predict_masks(
            self,
            image_embeddings: torch.Tensor,  # (1,256,64,64)
            image_pe: torch.Tensor,  # (1,256,64,64)  # position embeddings
            sparse_prompt_embeddings: torch.Tensor,  # (1,0,256)
            dense_prompt_embeddings: torch.Tensor,  # (1,256,64,64)
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Predicts masks. See 'forward' for more details."""
        # Concatenate output tokens
        output_tokens = torch.cat([self.iou_token.weight, self.mask_tokens.weight], dim=0)  # (1,256)+(4,256)=(5,256)
        output_tokens = output_tokens.unsqueeze(0).expand(sparse_prompt_embeddings.size(0), -1, -1)  # (5,256) -> (1,5,256)
        tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)  # (1,5,256)

        # Expand per-image data in batch direction to be per-mask
        src = torch.repeat_interleave(image_embeddings, tokens.shape[0], dim=0)  # (1,256,64,64)
        src = src + dense_prompt_embeddings
        pos_src = torch.repeat_interleave(image_pe, tokens.shape[0], dim=0)  # (1,256,64,64)
        b, c, h, w = src.shape  # 1,256,64,64

        # Run the transformer
        hs, src = self.transformer(src, pos_src, tokens)  # hs (1,5,256) src (1,4096,256)
        iou_token_out = hs[:, 0, :]  # (1,256)
        mask_tokens_out = hs[:, 1: (1 + self.num_mask_tokens), :]  # (1,4,256)

        # Upscale mask embeddings and predict masks using the mask tokens
        src = src.transpose(1, 2).view(b, c, h, w)  # (1,256,64,64)
        upscaled_embedding = self.output_upscaling(src)  # (1,32,256,256)
        upscaled_embedding_adaptor = self.output_upscaling_adaptor(
            src)  # (1,32,256,256) # new!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
        hyper_in_list: List[torch.Tensor] = []  # 4* (1,32)
        for i in range(self.num_mask_tokens):
            hyper_in_list.append(self.output_hypernetworks_mlps[i](mask_tokens_out[:, i, :]))
        hyper_in = torch.stack(hyper_in_list, dim=1)  # (1,4,32)
        b, c, h, w = upscaled_embedding.shape  # (1,32,256,256)

        ##### add MLP by pxy 230706 ######

        # scr_mlp_module = MLP(256,256,256,3)

        upscaled_embedding_concat = torch.cat([upscaled_embedding, upscaled_embedding_adaptor], dim=1)
        # new!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

        cls_upscaled_embedding = self.cls_upscaling(upscaled_embedding_concat)
        # new!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

        # cls_upscaled_embedding = self.cls_upscaling(upscaled_embedding)#原来的
        masks = (hyper_in @ cls_upscaled_embedding.view(b, c, 5 * h * w)).view(b, self.num_mask_tokens, -1, h, w)
        # (1,4,256,256)

        # masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w) #(1,4,256,256)#原来的

        # Generate mask quality predictions
        iou_pred = self.iou_prediction_head(iou_token_out)  # (1,4)

        return masks, iou_pred


# Lightly adapted from
# https://github.com/facebookresearch/MaskFormer/blob/main/mask_former/modeling/transformer/transformer_predictor.py # noqa
class MLP(nn.Module):
    def __init__(
            self,
            input_dim: int,
            hidden_dim: int,
            output_dim: int,
            num_layers: int,
            sigmoid_output: bool = False,
    ) -> None:
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(
            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
        )
        self.sigmoid_output = sigmoid_output

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        if self.sigmoid_output:
            x = F.sigmoid(x)
        return x
